# encoding:gbk
import pandas as pd
import logging
import sys
from datetime import datetime
import sqlite3
import time

class SQLiteDB:
    def __init__(self, db_name):
        """初始化数据库连接"""
        self.db_name = db_name
        self.conn = None
        self.cursor = None

    def connect(self):
        """连接到SQLite数据库"""
        self.conn = sqlite3.connect(self.db_name)
        self.cursor = self.conn.cursor()
        print(f"Connected to {self.db_name} database.")

    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.commit()
            self.conn.close()
            print(f"Connection to {self.db_name} closed.")

    def create_table(self, table_name, columns):
        """创建表格，传入表名和列的定义"""
        columns_str = ', '.join([f"{col} {dtype}" for col, dtype in columns.items()])
        query = f"CREATE TABLE IF NOT EXISTS {table_name} ({columns_str});"
        self.cursor.execute(query)
        print(f"Table '{table_name}' created or already exists.")

    def insert_data(self, table_name, data):
        """插入数据，data是一个字典形式的键值对"""
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['?' for _ in data])
        query = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"
        self.cursor.execute(query, tuple(data.values()))
        print(f"Data inserted into '{table_name}'.")

    def fetch_all(self, query, params=()):
        """执行查询，返回所有结果"""
        self.cursor.execute(query, params)
        return self.cursor.fetchall()

    def fetch_one(self, query, params=()):
        """执行查询，返回一个结果"""
        self.cursor.execute(query, params)
        return self.cursor.fetchone()

    def update_data(self, table_name, data, condition):
        """更新数据，data是一个字典，condition是更新条件"""
        set_clause = ', '.join([f"{key} = ?" for key in data])
        query = f"UPDATE {table_name} SET {set_clause} WHERE {condition}"
        self.cursor.execute(query, tuple(data.values()))
        print(f"Data in '{table_name}' updated.")

    def delete_data(self, table_name, condition):
        """删除数据，condition是删除条件"""
        query = f"DELETE FROM {table_name} WHERE {condition}"
        self.cursor.execute(query)
        print(f"Data from '{table_name}' deleted.")

def init(ContextInfo):
    """Initialize context parameters"""
    ContextInfo.accID = 'xxxxxx'
    ContextInfo.stock_symbols = ['513060.SH', '513090.SH', '510900.SH', '513980.SH', '164824.SZ']  # Changed to a list
    ContextInfo.db_file = r'C:\logs\qmt.sqlite'
    ContextInfo.order_list = []
    ContextInfo.threshold = 1000
    ContextInfo.turning_point_threshold = 0.003
    ContextInfo.drop_threshold = 0.99    # Turning point grid: buy on 1% drop
    ContextInfo.drop_threshold_1 = 0.995 # Normal grid
    ContextInfo.up_threshold = 1.01      # Turning point grid: sell on 1% rise
    ContextInfo.up_threshold_1 = 1.005   # Normal grid
    ContextInfo.start = '20250227 09:30:00'
    ContextInfo.end = '20250228 15:00:00'
    # Initialize dictionaries for each stock
    ContextInfo.base_prices = {symbol: get_base_price(ContextInfo, symbol) for symbol in ContextInfo.stock_symbols}
    ContextInfo.min_prices = {symbol: ContextInfo.base_prices[symbol] for symbol in ContextInfo.stock_symbols}
    ContextInfo.max_prices = {symbol: ContextInfo.base_prices[symbol] for symbol in ContextInfo.stock_symbols}
    ContextInfo.net_grid_data = {symbol: get_net_grid_data(ContextInfo, symbol) for symbol in ContextInfo.stock_symbols}
    ContextInfo.flag_dealed = {symbol: True for symbol in ContextInfo.stock_symbols}

def get_base_price(ContextInfo, stock_symbol):
    """Get base price for a specific stock"""
    db = SQLiteDB(ContextInfo.db_file)
    db.connect()
    query = "SELECT price FROM base_price WHERE stock = ?"
    result = db.fetch_one(query, (stock_symbol,))
    db.close()
    return float(result[0]) if result else None

def update_base_price(ContextInfo, stock_symbol):
    """Update base price for a specific stock"""
    orders = get_trade_detail_data(ContextInfo.accID, 'STOCK', 'ORDER')
    price_list = []
    for obj in orders:
        if '.'.join([obj.m_strInstrumentID, obj.m_strExchangeID]) == stock_symbol:
            price_list.append(obj.m_dTradedPrice)
    if price_list:
        ContextInfo.base_prices[stock_symbol] = price_list[-1]
        ContextInfo.min_prices[stock_symbol] = price_list[-1]
        ContextInfo.max_prices[stock_symbol] = price_list[-1]
        update_logs(ContextInfo, stock_symbol, f"更新基准价: {ContextInfo.base_prices[stock_symbol]}")
        
        # Save to database
        db = SQLiteDB(ContextInfo.db_file)
        db.connect()
        update_data = {'price': ContextInfo.base_prices[stock_symbol]}
        condition = f"stock = '{stock_symbol}'"
        db.update_data('base_price', update_data, condition)
        db.close()

def update_logs(ContextInfo, stock: str, data: str):
    """Save logs to database"""
    db = SQLiteDB(ContextInfo.db_file)
    db.connect()
    db.insert_data('logs', {'stock': stock, 'datetime': get_bartime(ContextInfo), 'log': data})
    db.close()

def get_net_grid_data(ContextInfo, stock_symbol):
    """Get grid data for a specific stock"""
    db = SQLiteDB(ContextInfo.db_file)
    db.connect()
    query = f"SELECT * FROM '{stock_symbol}'"
    result = db.fetch_all(query)
    db.close()
    return result

def turning_point_trigger(ContextInfo, stock_symbols):
    """Trigger trading based on turning points for a list of stocks"""
    if not isinstance(stock_symbols, list):
        stock_symbols = [stock_symbols]  # 确保输入为列表

    target = []  # 存储满足条件的股票代码

    for stock_symbol in stock_symbols:
        tick_data = ContextInfo.get_full_tick([stock_symbol])
        current_price = round(tick_data[stock_symbol]['lastPrice'], 3)
        base_price = ContextInfo.base_prices[stock_symbol]
        min_price = ContextInfo.min_prices[stock_symbol]
        max_price = ContextInfo.max_prices[stock_symbol]

        # 下跌趋势 Downtrend
        if round((current_price / base_price), 3) < ContextInfo.drop_threshold:
            update_logs(ContextInfo, stock_symbol, f"拐点触发|下跌趋势 - 最新价:{current_price}, 基准价:{base_price}, 最小价:{min_price}")
            if current_price < min_price:
                ContextInfo.min_prices[stock_symbol] = current_price
                update_logs(ContextInfo, stock_symbol, f"拐点触发|持续下跌 - 更新最低价: {ContextInfo.min_prices[stock_symbol]}")
            elif current_price > min_price * (1 + ContextInfo.turning_point_threshold):
                update_logs(ContextInfo, stock_symbol, f"拐点触发|下跌反弹 - 最新价: {current_price} > 最低价: {min_price} + 最低价 * {ContextInfo.turning_point_threshold}")
                target.append(stock_symbol)

        # 上涨趋势
        if round((current_price / base_price), 3) > ContextInfo.up_threshold:
            update_logs(ContextInfo, stock_symbol, f"拐点触发|上涨趋势 - 最新价:{current_price}, 基准价:{base_price}, 最高价:{max_price}")
            if current_price > max_price:
                ContextInfo.max_prices[stock_symbol] = current_price
                update_logs(ContextInfo, stock_symbol, f"拐点触发|持续上涨 - 更新最高价: {ContextInfo.max_prices[stock_symbol]}")
            elif current_price < max_price * (1 - ContextInfo.turning_point_threshold):
                update_logs(ContextInfo, stock_symbol, f"拐点触发|上涨回调 - 最新价:{current_price} < 最高价:{max_price} - 最高价 * {ContextInfo.turning_point_threshold}")
                target.append(stock_symbol)

    # Execute trade for all stocks in target list
    if target:
        exec_trade(ContextInfo, target)

def exec_trade(ContextInfo, stock_symbols:list):
    """Execute trade logic for a list of stocks"""
    for stock_symbol in stock_symbols:
        exec_cancel(ContextInfo, stock_symbol)
        tick_data = ContextInfo.get_full_tick([stock_symbol])
        current_price = round(tick_data[stock_symbol]['lastPrice'], 3)
        net_grid_data = ContextInfo.net_grid_data[stock_symbol]
        excel_quantity = [item[1] for item in net_grid_data if item[0] == current_price][0]
        position_data = get_holdings(ContextInfo.accID, 'STOCK')
        current_quantity = position_data.get(stock_symbol, {'持仓量': 0})['持仓量']
        update_logs(ContextInfo, stock_symbol, f"交易: 当前持仓:{current_quantity}, 表格持仓:{excel_quantity}")
        print(f"-- 交易标的：({stock_symbol})， 当前持仓:{current_quantity}, 表格持仓:{excel_quantity}")

        # Buy if Excel qty > Current qty
        if excel_quantity - current_quantity > ContextInfo.threshold:
            account_info = get_account(ContextInfo, ContextInfo.accID, 'STOCK')
            available_cash = account_info.get('可用金额', 0)
            trade_cost = current_price * (excel_quantity - current_quantity)
            if available_cash >= trade_cost:
                stk_price = tick_data[stock_symbol]['askPrice'][1]
                stk_amount = round((excel_quantity - current_quantity) / 100) * 100
                if ContextInfo.flag_dealed[stock_symbol]:
                    passorder(23, 1101, ContextInfo.accID, stock_symbol, 11, stk_price, stk_amount, 1, ContextInfo)
                    ContextInfo.flag_dealed[stock_symbol] = False
                update_logs(ContextInfo, stock_symbol, f"交易|买入 - 委托价:{stk_price}, 委托数量:{stk_amount} (Excel {excel_quantity} - Current {current_quantity})")
                print(f"-- 交易|买入： ({stock_symbol}) --\n委托价 {stk_price}, 委托数量 {stk_amount} (Excel {excel_quantity} - Current {current_quantity})")
            else:
                update_logs(ContextInfo, stock_symbol, f"交易|买入失败 - 余额不足：可用余额:{available_cash}, 交易金额:{trade_cost}")
                print(f"-- 交易|买入失败： ({stock_symbol}) --\n余额不足：可用余额:{available_cash}, 交易金额:{trade_cost}")

        # Sell if Current qty > Excel qty
        if current_quantity - excel_quantity > ContextInfo.threshold:
            stk_price = tick_data[stock_symbol]['bidPrice'][1]
            stk_amount = round((current_quantity - excel_quantity) / 100) * 100
            if ContextInfo.flag_dealed[stock_symbol]:
                passorder(24, 1101, ContextInfo.accID, stock_symbol, 11, stk_price, stk_amount, 1, ContextInfo)
                ContextInfo.flag_dealed[stock_symbol] = False
            update_logs(ContextInfo, stock_symbol, f"交易|卖出 - 委托价:{stk_price}, 委托数量:{stk_amount} (Current {current_quantity} - Excel {excel_quantity})")
            print(f"-- 交易|卖出： ({stock_symbol}) --\n委托价:{stk_price}, 委托数量:{stk_amount} (Current {current_quantity} - Excel {excel_quantity})")

def get_account(ContextInfo, accountid, datatype):
    """Get account data (placeholder, replace with actual implementation)"""
    accounts = get_trade_detail_data(accountid, datatype, 'account')
    result = {}
    for dt in accounts:
        result['总资产'] = dt.m_dBalance
        result['净资产'] = dt.m_dAssureAsset
        result['总市值'] = dt.m_dInstrumentValue
        result['总负债'] = dt.m_dTotalDebit
        result['可用金额'] = dt.m_dAvailable
        result['盈亏'] = dt.m_dPositionProfit
    return result

def get_holdings(accountid, datatype):
    """Get holdings data"""
    PositionInfo_dict = {}
    resultlist = get_trade_detail_data(accountid, datatype, 'POSITION')
    for obj in resultlist:
        stock_symbol = obj.m_strInstrumentID + '.' + obj.m_strExchangeID
        PositionInfo_dict[stock_symbol] = {
            '持仓量': obj.m_nVolume,
            '持仓成本': obj.m_dOpenPrice,
            '浮动盈亏': obj.m_dFloatProfit,
            '可用余额': obj.m_nCanUseVolume,
            '成交日期': obj.m_strOpenDate
        }
    return PositionInfo_dict

def exec_cancel(ContextInfo, stock_symbol):
    """Cancel unexecuted orders for a specific stock"""
    orders = get_trade_detail_data(ContextInfo.accID, 'STOCK', 'ORDER')
    for obj in orders:
        if obj.m_strInstrumentID == stock_symbol[:6]:
            if obj.m_strOrderSysID not in ContextInfo.order_list and obj.m_nVolumeTotal > 0:
                cancel(obj.m_strOrderSysID, ContextInfo.accID, 'STOCK', ContextInfo)
                update_logs(ContextInfo, stock_symbol, f"委托状态|撤单, 委托ID：{obj.m_strOrderSysID}")
                print(f"委托状态|撤单, 委托ID：{obj.m_strOrderSysID}")
            else:
                update_base_price(ContextInfo, stock_symbol)
                ContextInfo.flag_dealed[stock_symbol] = True
                update_logs(ContextInfo, stock_symbol, f"委托状态|成交, 委托ID {obj.m_strOrderSysID}")
                print(f"委托状态|成交, 委托ID {obj.m_strOrderSysID}")
                ContextInfo.order_list.append(obj.m_strOrderSysID)
    print(ContextInfo.order_list)

def get_bartime(ContextInfo, length=19):
    """Get current bar time"""
    bar_time = timetag_to_datetime(ContextInfo.get_bar_timetag(ContextInfo.barpos), '%Y-%m-%d %H:%M:%S')
    return bar_time[:length]

def main(ContextInfo):
    """Handle bar data for all stocks"""
    bar_time = timetag_to_datetime(ContextInfo.get_bar_timetag(ContextInfo.barpos), '%Y%m%d%H%M%S')
    target_normal = []
    target_guaidian = []
    if bar_time[-6:] >= '093000':  # Start after 9:30
        for stock_symbol in ContextInfo.stock_symbols:
            tick_data = ContextInfo.get_full_tick([stock_symbol])
            last_price = tick_data[stock_symbol]['lastPrice']
            last_close = tick_data[stock_symbol]['lastClose']
            if 0.99 < last_price/last_close < 1.01:
                if round((last_price/ContextInfo.base_prices[stock_symbol]),3) <= ContextInfo.drop_threshold_1 or round((last_price/ContextInfo.base_prices[stock_symbol]),3) > ContextInfo.up_threshold_1:
                    target_normal.append(stock_symbol)
            else:
                if round((last_price/ContextInfo.base_prices[stock_symbol]),3) <= ContextInfo.drop_threshold or round((last_price/ContextInfo.base_prices[stock_symbol]),3) > ContextInfo.up_threshold:
                    target_guaidian.append(stock_symbol)
        if len(target_normal) > 0:
            exec_trade(ContextInfo, target_normal)
            print('普通触发')
        if len(target_guaidian) > 0:
            turning_point_trigger(ContextInfo, target_guaidian)
            print('拐点触发')
    else:
        print('非交易时间')

def handlebar(ContextInfo):
    flag = 1
    if flag == 1:
        main(ContextInfo)
    elif flag == 2:
        test(ContextInfo)

def test(ContextInfo):
    stock_symbol = '164824.SZ'
    position_data = get_holdings(ContextInfo.accID, 'STOCK')
    current_quantity = position_data
    print(current_quantity)















